#include <stdio.h>
#include <cuda_runtime.h>

// 向量加法CUDA内核
__global__ void add(int *a, int *b, int *c, int n) {
	int index = threadIdx.x + blockIdx.x * blockDim.x;
	if (index < n) {
		c[index] = a[index] + b[index];
	}
}

int main() {
	int n = 1024; // 向量大小
	size_t bytes = n * sizeof(int);
	int *h_a, *h_b, *h_c; // 主机端向量
	int *d_a, *d_b, *d_c; // 设备端向量
	int *h_pinned_a, *h_pinned_b; // 锁页内存指针

	// 分配主机内存
	h_a = (int *)malloc(bytes);
	h_b = (int *)malloc(bytes);
	h_c = (int *)malloc(bytes);

	// 初始化主机端向量
	for (int i = 0; i < n; i++) {
		h_a[i] = i;
		h_b[i] = 2 * i;
	}

	// 分配锁页内存
	cudaHostAlloc(&h_pinned_a, bytes, cudaHostAllocPortable);
	cudaHostAlloc(&h_pinned_b, bytes, cudaHostAllocPortable);

	// 从主机复制到锁页内存
	cudaMemcpy(h_pinned_a, h_a, bytes, cudaMemcpyHostToHost);
	cudaMemcpy(h_pinned_b, h_b, bytes, cudaMemcpyHostToHost);

	// 分配设备内存
	cudaMalloc(&d_a, bytes);
	cudaMalloc(&d_b, bytes);
	cudaMalloc(&d_c, bytes);

	// 创建两个流
	cudaStream_t streamCopy, streamKernel;
	cudaStreamCreate(&streamCopy);
	cudaStreamCreate(&streamKernel);

	// 从锁页内存异步复制到设备，使用streamCopy
	cudaMemcpyAsync(d_a, h_pinned_a, bytes, cudaMemcpyHostToDevice, streamCopy);
	cudaMemcpyAsync(d_b, h_pinned_b, bytes, cudaMemcpyHostToDevice, streamCopy);

	// 执行内核，使用streamKernel
	int threadsPerBlock = 256;
	int blocksPerGrid = (n + threadsPerBlock - 1) / threadsPerBlock;
	add<<<blocksPerGrid, threadsPerBlock, 0, streamKernel>>>(d_a, d_b, d_c, n);

	// 等待内核执行完成
	cudaStreamSynchronize(streamKernel);

	// 从设备异步复制回锁页内存，使用streamCopy
	cudaMemcpyAsync(h_pinned_a, d_c, bytes, cudaMemcpyDeviceToHost, streamCopy);

	// 等待数据传输完成
	cudaStreamSynchronize(streamCopy);

	// 从锁页内存复制回主机内存
	cudaMemcpy(h_c, h_pinned_a, bytes, cudaMemcpyHostToHost); // 注意这里应该是h_pinned_a，而不是h_pinned_b

	// 检查结果
	int i;
	for (i = 0; i < n; i++) {
		if (h_c[i] != h_a[i] + h_b[i]) {
			printf("Error: invalid result at %d\n", i);
			break;
		}
	}

	if (i >= n) {
		printf("Vector addition is correct.\n");
	}

	// 清理资源
	free(h_a);
	free(h_b);
	free(h_c);
	cudaFreeHost(h_pinned_a);
	cudaFreeHost(h_pinned_b);
	cudaFree(d_a);
	cudaFree(d_b);
	cudaFree(d_c);
	cudaStreamDestroy(streamCopy);
	cudaStreamDestroy(streamKernel);


	return 0;
}
